"""
热传导方程求解器

该模块实现了一维圆柱热传导方程的求解,用于计算燃料棒包壳的温度分布。
基于论文中的公式：
ρc*Cp*∂T/∂t = (1/r)*∂/∂r(r*λc*∂T/∂r) + Qin - Qv - Ql
"""

import numpy as np
from typing import Dict, Tuple, Optional
from scipy.sparse import diags
from scipy.sparse.linalg import spsolve

class HeatConductionSolver:
    """一维圆柱热传导方程求解器"""
    
    def __init__(self, geometry: Dict):
        """
        初始化求解器
        
        参数:
            geometry: 几何参数字典,包含燃料棒直径、网格数等
        """
        self.geometry = geometry
        self.n_radial = 20  # 径向网格数
        
        # 计算网格参数
        self.r_inner = 0  # 燃料棒中心
        self.r_outer = geometry['rod_diameter'] / 2  # 燃料棒外径
        self.dr = (self.r_outer - self.r_inner) / self.n_radial
        
        # 初始化温度场
        self.temperature = np.zeros(self.n_radial)
        
        # 材料属性
        self.rho_c = 6500.0  # 包壳密度 (kg/m³)
        self.cp_c = 330.0    # 包壳比热容 (J/kg·K)
        self.lambda_c = 16.0 # 包壳导热系数 (W/m·K)
        
    def solve_steady_state(self, q_in: float, h_v: float, h_l: float,
                          t_v: float, t_l: float, t_wall: float) -> np.ndarray:
        """
        求解稳态热传导方程
        
        参数:
            q_in: 进入包壳的热流密度 (W/m²)
            h_v: 气相对流传热系数 (W/m²·K)
            h_l: 液相对流传热系数 (W/m²·K)
            t_v: 气相温度 (K)
            t_l: 液相温度 (K)
            t_wall: 壁面温度初值 (K)
            
        返回:
            径向温度分布
        """
        # 构建系数矩阵
        r = np.linspace(self.r_inner + self.dr/2, 
                       self.r_outer - self.dr/2, self.n_radial)
        
        # 主对角线系数
        main_diag = -2 * np.ones(self.n_radial)
        main_diag[0] = -1  # 中心边界条件
        main_diag[-1] = -1 - self.dr * (h_v + h_l) / self.lambda_c
        
        # 上下对角线系数
        upper_diag = np.ones(self.n_radial-1) * (1 + self.dr/(2*r[1:]))
        lower_diag = np.ones(self.n_radial-1) * (1 - self.dr/(2*r[:-1]))
        
        # 构建三对角矩阵
        A = diags([lower_diag, main_diag, upper_diag], [-1, 0, 1]).tocsr()
        
        # 构建右端向量
        b = np.zeros(self.n_radial)
        b[0] = 0  # 中心对称边界条件
        b[-1] = -self.dr * (h_v*t_v + h_l*t_l) / self.lambda_c - \
                self.dr * q_in / self.lambda_c
        
        # 求解线性方程组
        self.temperature = spsolve(A, b)
        
        return self.temperature
        
    def solve_transient(self, dt: float, q_in: float, h_v: float, h_l: float,
                       t_v: float, t_l: float) -> np.ndarray:
        """
        求解瞬态热传导方程
        
        参数:
            dt: 时间步长 (s)
            q_in: 进入包壳的热流密度 (W/m²)
            h_v: 气相对流传热系数 (W/m²·K)
            h_l: 液相对流传热系数 (W/m²·K)
            t_v: 气相温度 (K)
            t_l: 液相温度 (K)
            
        返回:
            新的径向温度分布
        """
        r = np.linspace(self.r_inner + self.dr/2, 
                       self.r_outer - self.dr/2, self.n_radial)
        
        # 计算扩散系数
        alpha = self.lambda_c / (self.rho_c * self.cp_c)
        
        # 构建系数矩阵
        main_diag = np.ones(self.n_radial) * (1 + 2*alpha*dt/self.dr**2)
        main_diag[0] = 1 + alpha*dt/self.dr**2
        main_diag[-1] = 1 + alpha*dt/self.dr**2 + \
                       dt*(h_v + h_l)/(self.rho_c * self.cp_c)
        
        upper_diag = np.ones(self.n_radial-1) * \
                    (-alpha*dt/self.dr**2 * (1 + self.dr/(2*r[1:])))
        lower_diag = np.ones(self.n_radial-1) * \
                    (-alpha*dt/self.dr**2 * (1 - self.dr/(2*r[:-1])))
        
        A = diags([lower_diag, main_diag, upper_diag], [-1, 0, 1]).tocsr()
        
        # 构建右端向量
        b = np.copy(self.temperature)
        b[0] = self.temperature[0]  # 中心对称边界条件
        b[-1] = self.temperature[-1] + \
                dt/(self.rho_c * self.cp_c) * \
                (q_in + h_v*t_v + h_l*t_l)
        
        # 求解线性方程组
        self.temperature = spsolve(A, b)
        
        return self.temperature
        
    def get_wall_temperature(self) -> float:
        """获取壁面温度"""
        return self.temperature[-1]
        
    def get_heat_flux(self) -> Tuple[float, float]:
        """
        计算传向气相和液相的热流密度
        
        返回:
            (传向气相的热流密度, 传向液相的热流密度)
        """
        # 使用傅里叶导热定律计算壁面热流
        q_wall = -self.lambda_c * \
                 (self.temperature[-1] - self.temperature[-2]) / self.dr
        
        # 假设热流均匀分配给气相和液相
        return q_wall/2, q_wall/2 